#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
from multiprocessing import Queue, Pool
import multiprocessing
import time
import fire
import os

MONGODB_HOST = "10.10.117.188"
MONGODB_PORT = 27017
MONGODB_DB = "hanvcf"
MONGODB_USER = "han"
MONGODB_PASSWORD = "han2019"
MONGODB_COLLECTION_GENOTYPES = "genotypes_20190801"
MONGODB_COLLECTION_SNVS = "snvs_20190801"


# # 初始写入的代码:
# def write(col: collection):
#     for i in range(100):
#         col.update_one({'test_id': i}, {'$set': {'test_id': i}}, upsert=True)
#
#
# # 读入ID的进程代码:
# def read_id(ip: str, db: str, col: str, q: Queue):
#     print("Read session start.")
#     client = MongoClient(ip, 27017)
#     database = client[db]
#     collection = database[col]
#     for result in collection.find({}, {"_id": 1}):
#         q.put(result["_id"])
#     for i in range(10):
#         q.put("End")
#     client.close()
#     print("Read session complete.")
#
#
# # 修改ID的进程代码:
# def write_id(ip: str, db: str, col: str, q: Queue):
#     print("Update session start.")
#     client = MongoClient(ip, 27017)
#     database = client[db]
#     collection = database[col]
#     while True:
#         if not q.empty():
#             value = q.get()
#             if str(value) != "End":
#                 collection.update_one({"_id": value}, {"$set": {"test_id": random.random()}})
#             else:
#                 client.close()
#                 print("Update session completed.")
#                 break


def add_to_genotypes_queue(q: Queue, q_max: int):
    client_genotypes = MongoClient(MONGODB_HOST, MONGODB_PORT)
    database_genotypes = client_genotypes[MONGODB_DB]
    database_genotypes.authenticate(MONGODB_USER, MONGODB_PASSWORD)
    collection_genotypes = database_genotypes[MONGODB_COLLECTION_GENOTYPES]

    for result in collection_genotypes.find({}, {"_id": 1}, no_cursor_timeout=True):
        while q.qsize() >= (q_max * 0.8):
            time.sleep(5)
        q.put(result["_id"])

    client_genotypes.close()


def add_to_snvs_queue(q: Queue, q_max: int):
    client_snvs = MongoClient(MONGODB_HOST, MONGODB_PORT)
    database_snvs = client_snvs[MONGODB_DB]
    database_snvs.authenticate(MONGODB_USER, MONGODB_PASSWORD)
    collection_snvs = database_snvs[MONGODB_COLLECTION_SNVS]

    for result in collection_snvs.find({}, {"_id": 1}, no_cursor_timeout=True):
        while q.qsize() >= (q_max * 0.8):
            time.sleep(5)
        q.put(result["_id"])

    client_snvs.close()


# 根据Allele Frequency（geno）计算SNV发生率的函数
def get_possibility(ref: str, alt: str, ances: str, geno: dict):
    total_individual_number = 0
    total_alt_number = 0

    if ances == alt:
        allele = ref
    else:
        allele = alt

    for item in geno.keys():
        total_individual_number += geno[item]
        count_alt = item.count(allele)
        total_alt_number += count_alt * geno[item]

    total_individual_number = total_individual_number * 2

    return [total_alt_number, total_individual_number]


# 在genotypes表中更新人群（pop_aaf）和大区(region_aaf)的脚本
def update_pop_region_aaf(q: Queue):
    client_genotypes = MongoClient(MONGODB_HOST, MONGODB_PORT)
    database_genotypes = client_genotypes[MONGODB_DB]
    database_genotypes.authenticate(MONGODB_USER, MONGODB_PASSWORD)
    collection_genotypes = database_genotypes[MONGODB_COLLECTION_GENOTYPES]

    count = 0
    handled = 0
    existed = 0
    print("Update process " + str(os.getpid()) + " completed.")
    while True:
        if not q.empty():
            q_value = q.get()
            if str(q_value) != "End":
                result = collection_genotypes.find({"_id": q_value})
                if "pop_aaf" not in result:
                    if "geno" in result:
                        pop_aaf = get_possibility(result["ref"], result["alt"], result["ances"], result["geno"])
                        collection_genotypes.update_one({"_id": result["_id"]}, {
                            "$set": {"pop_aaf": {"alt_number": pop_aaf[0], "individual_number": pop_aaf[1]}}})
                    else:
                        collection_genotypes.update_one({"_id": result["_id"]}, {
                            "$set": {"pop_aaf": {"alt_number": 0, "individual_number": 0}}})
                    handled += 1
                else:
                    existed += 1
                if "region_aaf" not in result:
                    region_aaf = [0, 0]
                    for region_result in collection_genotypes.find(
                            {"chrom": result["chrom"], "pos": result["pos"], "ref": result["ref"],
                             "alt": result["alt"],
                             "genetic_region": result["genetic_region"]}):
                        if "geno" in region_result:
                            cpop_aaf = get_possibility(region_result["ref"], region_result["alt"],
                                                       region_result["ances"],
                                                       region_result["geno"])
                            region_aaf[0] += cpop_aaf[0]
                            region_aaf[1] += cpop_aaf[1]
                    collection_genotypes.update_many(
                        {"chrom": result["chrom"], "pos": result["pos"], "ref": result["ref"], "alt": result["alt"],
                         "genetic_region": result["genetic_region"]},
                        {"$set": {"region_aaf": {"alt_number": region_aaf[0], "individual_number": region_aaf[1]}}})
                count += 1
                if count % 100000 == 0:
                    print("Update process " + str(os.getpid()) + " : " + str(count) + " records processed, " + str(
                        handled) + " records added, " + str(existed) + " records skipped.")
            else:
                client_genotypes.close()
                print("Update process " + str(os.getpid()) + " completed.")
                break


# 在snv表中更新SNV发生频率（snv_aaf）的脚本
def update_snv_aaf(q: Queue):
    client_snvs = MongoClient(MONGODB_HOST, MONGODB_PORT)
    database_snvs = client_snvs[MONGODB_DB]
    database_snvs.authenticate(MONGODB_USER, MONGODB_PASSWORD)
    collection_genotypes = database_snvs[MONGODB_COLLECTION_GENOTYPES]
    collection_snvs = database_snvs[MONGODB_COLLECTION_SNVS]

    count = 0
    handled = 0
    existed = 0
    print("Update process " + str(os.getpid()) + " completed.")
    while True:
        if not q.empty():
            q_value = q.get()
            if str(q_value) != "End":
                result = collection_snvs.find({"_id": q_value})
                if "snv_aaf" not in result:
                    snv_aaf = [0, 0]
                    for pop_result in collection_genotypes.find(
                            {"chrom": result["chrom"], "pos": result["pos"], "ref": result["ref"],
                             "alt": result["alt"]}):
                        if "geno" in pop_result:
                            cpop_aaf = get_possibility(pop_result["ref"], pop_result["alt"], pop_result["ances"],
                                                       pop_result["geno"])
                            snv_aaf[0] += cpop_aaf[0]
                            snv_aaf[1] += cpop_aaf[1]
                    collection_snvs.update_one({"_id": result["_id"]}, {
                        "$set": {"snv_aaf": {"alt_number": snv_aaf[0], "individual_number": snv_aaf[1]}}})
                    handled += 1
                else:
                    existed += 1
                count += 1
                if count % 10000 == 0:
                    print("Update process " + str(os.getpid()) + " : " + str(count) + " records processed, " + str(
                        handled) + " records added, " + str(existed) + " records skipped.")
            else:
                client_snvs.close()
                print("Update process " + str(os.getpid()) + " completed.")
                break


def update_pop_region_aaf_main(qsize: int):
    manager = multiprocessing.Manager()
    q = manager.Queue(qsize)
    lock = manager.Lock()
    p = Pool()
    p.apply_async(add_to_genotypes_queue, args=(q, qsize))
    for i in range(multiprocessing.cpu_count() - 3):
        p.apply_async(update_pop_region_aaf, args=(q,))
    p.close()
    p.join()


def update_snvs_aaf_main(qsize: int):
    manager = multiprocessing.Manager()
    q = manager.Queue(qsize)
    lock = manager.Lock()
    p = Pool()
    p.apply_async(add_to_snvs_queue, args=(q, qsize))
    for i in range(multiprocessing.cpu_count() - 3):
        p.apply_async(update_snv_aaf, args=(q,))
    p.close()
    p.join()


# def UpdateRegion():
#     htfile = open("province_region.txt", "r")
#     province_region = {}
#     for line in htfile:
#         line = line.rstrip()
#         splitline = line.split("\t")
#         province_region[splitline[0]] = splitline[1]
#     htfile.close()
#
#     client_genotypes, collection_genotypes = InitMongoDB("localhost", 27017, "hanvcf", "genotypes")
#
#     for result in collection_genotypes.find():
#         collection_genotypes.update_one({"_id": result["_id"]},
#                                         {"$set": {"genetic_region": province_region[result["pop_name"]]}})
#
#     CloseMongoDB(client_genotypes)


if __name__ == '__main__':
    fire.Fire()
